from typing import List, Tuple, Dict, Optional, Set
import dbutil

class AreaItem:
    """区域项，存储编码和关键词"""
    def __init__(self, code: str, keyword: str):
        self.code = code
        self.keyword = keyword
        
    def __repr__(self):
        return f"AreaItem(code='{self.code}', keyword='{self.keyword}')"
    
    def __eq__(self, other):
        if isinstance(other, AreaItem):
            return self.code == other.code and self.keyword == other.keyword
        return False
    
    def __hash__(self):
        return hash((self.code, self.keyword))

class Node:
    """Aho-Corasick自动机的节点"""
    def __init__(self):
        self.goto: Dict[str, Node] = {}  # 转移表
        self.fail: Optional[Node] = None  # 失败指针
        self.output: Set[str] = set()    # 输出关键词集合

class Automaton:
    """自定义Aho-Corasick自动机实现多模匹配"""
    def __init__(self):
        self.root = Node()
    
    def add_word(self, word: str) -> None:
        """添加关键词到自动机"""
        node = self.root
        for char in word:
            if char not in node.goto:
                node.goto[char] = Node()
            node = node.goto[char]
        node.output.add(word)
    
    def make_automaton(self) -> None:
        """构建自动机的失败指针，完成自动机初始化"""
        queue = []
        
        # 初始化根节点的所有子节点的失败指针为根节点
        for node in self.root.goto.values():
            node.fail = self.root
            queue.append(node)
        
        # BFS构建失败指针
        while queue:
            current_node = queue.pop(0)
            
            # 处理当前节点的每个子节点
            for char, child in current_node.goto.items():
                # 找到失败指针
                fail_node = current_node.fail
                while fail_node is not None and char not in fail_node.goto:
                    fail_node = fail_node.fail
                
                child.fail = fail_node.goto[char] if fail_node and char in fail_node.goto else self.root
                # 合并输出
                child.output.update(child.fail.output)
                queue.append(child)
    
    def iter(self, text: str) -> List[Tuple[int, str]]:
        """在文本中查找所有匹配的关键词，返回(结束位置, 关键词)"""
        results = []
        current = self.root
        
        for i, char in enumerate(text):
            # 沿着失败链查找
            while current is not None and char not in current.goto:
                current = current.fail
            
            if current is None:
                current = self.root
            else:
                current = current.goto.get(char, self.root)
            
            # 收集所有匹配的关键词
            for word in current.output:
                results.append((i, word))
        
        return results

class AreaKeywords:
    """区域关键词匹配工具类"""
    def __init__(self):
        self.match_object: Optional[Automaton] = None  # 多模匹配对象
        self.keyword_map: Dict[str, List[AreaItem]] = {}  # 关键词到AreaItem的映射
    
    def load_keywords(self, level: int) -> int:
        """从数据库加载指定级别的区域关键词"""
        table_mapping = {
            2: 'area2_feature',
            3: 'area3_feature',
            4: 'area4_feature',
            5: 'area5_feature'
        }
        
        if level not in table_mapping:
            raise ValueError(f"不支持的级别: {level}，请使用2-5之间的级别")
            
        table = table_mapping[level]
        sql = f"select code, full_name, alias1, alias2, alias3, alias4 from {table}"
        rows = dbutil.query_rows(sql, ())
        
        self.match_object = Automaton()
        self.keyword_map = {}
        keyword_count = 0
        
        for row in rows:
            # 提取字段并处理空值和空格
            code = str(row.get('code', '')).strip()
            full_name = str(row.get('full_name', '')).strip()
            alias1 = str(row.get('alias1', '')).strip()
            alias2 = str(row.get('alias2', '')).strip()
            alias3 = str(row.get('alias3', '')).strip()
            alias4 = str(row.get('alias4', '')).strip()
            
            # 收集非空关键词
            keywords = []
            for kw in [full_name, alias1, alias2, alias3, alias4]:
                if kw:  # 只添加非空关键词
                    keywords.append(kw)
            
            # 为每个关键词添加到自动机
            for kw in keywords:
                if kw not in self.keyword_map:
                    self.keyword_map[kw] = []
                    self.match_object.add_word(kw)
                    keyword_count += 1
                self.keyword_map[kw].append(AreaItem(code, kw))
        
        # 构建自动机
        if keyword_count > 0:
            self.match_object.make_automaton()
            
        return keyword_count
    
    def match(self, address: str) -> List[AreaItem]:
        """匹配地址中的区域关键词"""
        if not self.match_object or not address:
            return []
            
        matched_items = set()
        
        # 多模匹配
        for item in self.match_object.iter(address):
            keyword = item[1]
            if keyword in self.keyword_map:
                for area_item in self.keyword_map[keyword]:
                    matched_items.add(area_item)
        
        return list(matched_items)

# 测试代码
def test_area_keywords():
    area_matcher = AreaKeywords()
    count = area_matcher.load_keywords(3)
    print(f"加载了 {count} 个关键词")
    
    address = "中共察布查尔锡伯自治县委员会党校"
    matched_items = area_matcher.match(address)
    
    print(f"地址 '{address}' 的匹配结果:")
    for item in matched_items:
        print(f"编码: {item.code}, 匹配关键词: {item.keyword}")

if __name__ == "__main__":
    test_area_keywords()
    